from fastapi import FastAPI
from fastapi.responses import RedirectResponse
from langserve import add_routes
from langchain_community.llms import Ollama
from langchain_core.prompts import ChatPromptTemplate
# 添加星火大模型路由
from langchain_community.llms import SparkLLM
from pyspark.sql import SparkSession 
from typing import List

# 创建FastAPI应用
app = FastAPI(
    title="LangChain Server",
    version="1.0",
    description="A simple API server using LangChain with Ollama"
)

@app.get("/")
async def redirect_root_to_docs():
    return RedirectResponse("/docs")

################################# 配置Ollama###################################
# llm = Ollama(
#     model="deepseek-r1:latest", 
#     temperature=0,
#     top_p=0.,
#     base_url="http://192.168.13.201:11434"
# )

# # 创建prompt模板
# prompt = ChatPromptTemplate.from_messages([
#     ("system", "你是一位世界级的技术文档写手，你的名字是贾维斯"),
#     ("user", "{input}")
# ])

# # 创建chain
# chain = prompt | llm
# 添加路由
# add_routes(
#     app,
#     chain,
#     path="/chat",  # API端点
# )
################################# 配置Ollama###################################


################################# 配置星火大模型###################################




# 配置星火大模型，设置websocket地址
# 配置星火大模型的websocket地址
spark_ws_url = "wss://spark-api.xf-yun.com/v4.1/chat"

spark_llm = SparkLLM(
    app_id="901f3979",
    api_key="b86297cdc8ba79316c6145fe48baeb77", 
    api_secret="ZGYyMmRlNGE1NGY0YjQ1ZWE1ZDQ4NWUx",
    temperature=0.7,
    # domain="generalv3.5", # 使用V3.5版本
    websocket_url=spark_ws_url
)

# 创建星火prompt模板
spark_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一位专业的AI助手,你的名字是星火"),
    ("user", "{input}")
])

# 创建星火chain
spark_chain = spark_prompt | spark_llm

# 添加星火路由
add_routes(
    app,
    spark_chain,
    path="/spark",  # 星火API端点
)

# 创建 Spark 会话
spark = SparkSession.builder \
    .appName("LangChain with Spark 4.0 Ultra") \
    .getOrCreate()

@app.post("/spark/process")
async def process_data(data: List[dict]):
    # 将输入数据转换为 DataFrame
    df = spark.createDataFrame(data)
    
    # 进行一些数据处理，例如计算总和
    result = df.agg({"value": "sum"}).collect()
    
    return {"result": result}

################################# 配置星火大模型###################################

# Edit this to add the chain you want to add
# add_routes(app, NotImplemented)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
